from typing import Any

import numpy as np

from ..dataclass.results import SingleResult, UnitResult
from ..pipeline.nameSpace import get_local_function
from ..pipeline.pipeData import NestedRangeModel, RangeModel
from ..pipeline.pipeline import PipeModule, Port, status_manage
from ..pipeline.portTypes import PortType, PortTypeHint


@status_manage
class ItemsSelector(PipeModule):
    """Define a selector to select multi-items from a list."""

    def __init__(
        self,
        mname: str | None = None,
        auto_run: bool = True,
        items_list: list[Any] | None = None,
        items_label: list[str] | None = None,
        single_selector: bool = True,
        auto_selected: list[Any] | None = None,
        bunddle_items: dict[str, list[Any]] | None = None,
        as_single_result: bool = False,
        single_unit_result: bool = False,
    ) -> None:
        """Initialize ItemsSelector object.

        Parameters
        ----------
        items_list : list[Any]
            The items to select from.

        items_label : list[str]
            The labels of the items which is used to show the users.

        single_selector : bool
            Whether to use a single selector or a multi-selector.

        auto_selected : list[Any]
            The items to select automatically.
            These items usually are not visible to users.

        bunddle_items : dict[str, list[Any]]
            Bunddle serveral items to a single item.
            key is the name of the item in the drop list,
            value is the output items.
            Exmaple:
            {
                "item1": ["item1_1", "item1_2", "item1_3"],
                "item2": ["item2_1", "item2_2", "item2_3"]
            }

        as_single_result : bool, default: False
            If True, the output will be a SingleResult object.
            If False, the output will be a general array.

        single_unit_result : bool, default: False
            Only available when as_single_result is True.
            If True, all the seleted items will be a value of a unit reuslt in a list type whiche means the length of the single result will always be 1.
            If False, each selected item will be a value of a unit result which means the length of the single result will equal to the length of the selected items.
        """
        super().__init__(mname, auto_run)

        self._ports_out = {
            "OutputItems": Port(
                [PortType.GeneralArray, PortType.SingleResult],
                pdoc="The selected items.",
            )
        }

        self.items_list = items_list
        if items_label:
            if len(items_label) != len(items_list):
                raise ValueError(
                    "The length of items_label must be equal to the length of items_list."
                )
        self.items_label = items_label or items_list
        self.single_selector = single_selector
        self.auto_selected = auto_selected
        self.bunddle_items = bunddle_items
        self.selected_items: list[Any] | None = None
        self.as_single_result = as_single_result
        self.single_unit_result = single_unit_result

    @property
    def OutputItems(self) -> PortTypeHint.GeneralArray | None:
        return self["OutputItems"]

    def set_cal_params(self, reset: bool = False) -> dict[str, RangeModel]:
        if self.single_selector:
            list_len = 1
        else:
            list_len = None

        values_range = {
            "selected_items": RangeModel(
                vtype=("list", "None"),
                title="可选项",
                choices=self.items_list,
                choices_label=self.items_label,
                list_len=list_len,
            )
        }
        return values_range

    def execute(self) -> PortTypeHint.GeneralArray | None:
        if self.selected_items is None or self.items_list is None:
            self._ports_out["OutputItems"].data = None
            return None

        selected_items = []
        if self.auto_selected:
            selected_items = self.auto_selected.copy()
        if self.bunddle_items is None:
            selected_items.extend(self.selected_items)
        else:
            for item in self.selected_items:
                if item in self.bunddle_items:
                    selected_items.extend(self.bunddle_items[item])
                else:
                    selected_items.append(item)

        if self.as_single_result:
            unit_result_list = []
            if self.single_unit_result:
                unit_result = UnitResult(
                    name="selected_items", title="选择的项", value=selected_items
                )
                unit_result_list.append(unit_result)
            else:
                for idx, item in enumerate(selected_items):
                    unit_result = UnitResult(
                        name=f"selected_item_{idx}", title=item, value=item
                    )
                    unit_result_list.append(unit_result)
            single_result = SingleResult(unit_result_list)

            self._ports_out["OutputItems"].data = single_result
            return single_result

        self._ports_out["OutputItems"].data = selected_items
        return selected_items


@status_manage
class AttributeCaculator(PipeModule):
    """This module is used to define a local function to calculate the values of attributes for another module."""

    def __init__(
        self,
        mname: str | None = None,
        auto_run: bool = True,
        local_functions_path: str | None = None,
        local_function_name: str | None = None,
        pipeline_local_functions_path_first: bool = True,
    ) -> None:
        """Initialize AttributeCaculator object.

        Parameters
        ----------
        local_functions_path : str | None, default: None
            The path to the local functions file.
            If it's not provided, the module will try to use the local functions file path in the pipeline.

        local_function_name : str | None, default: None
            The name of the class that defines the function.

        pipeline_local_functions_path_first : bool, default: True
            If True, the 'local_functions_path' attribute of the pipeline will be used to find the local functions path first.
            If False, the module's attributes 'local_functions_path' will be used to find the local functions path first.
        """
        super().__init__(mname, auto_run)

        self._ports_out = {
            "OutputAttributesValue": Port(
                PortType.Attributes, pdoc="The attributes value for child module."
            )
        }

        self.local_functions_path = local_functions_path
        self.local_function_name = local_function_name
        self.pipeline_local_functions_path_first = pipeline_local_functions_path_first

    @property
    def OutputAttributesValue(self) -> PortTypeHint.Attributes | None:
        return self["OutputAttributesValue"]

    def set_cal_params(self, reset: bool = True) -> None:
        return None

    def execute(self) -> PortTypeHint.Attributes | None:
        if self.local_function_name is None:
            self._ports_out["OutputAttributesValue"].data = None
            return None

        # Get the local function
        calculator = get_local_function(
            self, self.local_function_name, self.pipeline_local_functions_path_first
        )
        if calculator is None:
            self._ports_out["OutputAttributesValue"].data = None
            return None

        result = calculator.run_function()
        if result:
            if not isinstance(result, dict):
                raise ValueError(
                    f"The result of the function '{self.local_function_name}' (type: {type(result)}) is not a dict."
                )
            self._ports_out["OutputAttributesValue"].data = result
            return result

        self._ports_out["OutputAttributesValue"].data = None
        return None


@status_manage
class PythonCoder(PipeModule):
    """This module is used to run a local python script.

    This is a flexible module that can accept any type of input through General ports
    and output any type of data through General ports by running custom Python functions.
    """

    def __init__(
        self,
        mname: str | None = None,
        auto_run: bool = True,
        local_functions_path: str | None = None,
        local_function_name: str | None = None,
        range_model_function_name: str | None = None,
        pipeline_local_functions_path_first: bool = True,
        debug_mode: bool = False,
    ) -> None:
        """Initialize PythonCoder object.

        Parameters
        ----------
        local_functions_path : str | None, default: None
            The path to the local functions file.
            If it's not provided, the module will try to use the local functions file path in the pipeline.

        local_function_name : str | None, default: None
            The name of the class that defines the function.

        range_model_function_name : str | None, default: None
            The name of the function that returns range model.

        pipeline_local_functions_path_first : bool, default: True
            If True, the 'local_functions_path' attribute of the pipeline will be used to find the local functions path first.
            If False, the module's attributes 'local_functions_path' will be used to find the local functions path first.

        debug_mode : bool, default: False
            If True, provides detailed error messages including full tracebacks when loading local functions fails.
            If False, provides concise error messages for production use.
            Useful for debugging syntax errors, import errors, and other issues in local function files.

        Note
        ----
        The returned result from the local function must be a dict: key is the name of dynamic_ports_out, value is the port data.
        """
        super().__init__(mname, auto_run)

        # Use General port types to accept and output any kind of data
        # Use Dynamic Port to define any numbers of ports
        self.dynamic_ports_in_type = PortType.General
        self.dynamic_ports_out_type = PortType.General

        self.local_functions_path = local_functions_path
        self.local_function_name = local_function_name
        self.range_model_function_name = range_model_function_name
        self.pipeline_local_functions_path_first = pipeline_local_functions_path_first
        self.debug_mode = debug_mode

    def set_cal_params(
        self, reset: bool = True
    ) -> dict[str, RangeModel | NestedRangeModel] | None:
        if self.range_model_function_name:
            calculator = get_local_function(
                self,
                self.range_model_function_name,
                self.pipeline_local_functions_path_first,
                self.debug_mode,
            )
            if calculator is None:
                return None

            result = calculator.run_function(module=self, reset=reset)
            if result is None:
                return None
            if not isinstance(result, dict):
                raise ValueError(
                    f"The result of the function '{self.range_model_function_name}' (type: {type(result)}) is not a dict."
                )
            for key, value in result.items():
                if not isinstance(key, str) or not isinstance(
                    value, (RangeModel, NestedRangeModel)
                ):
                    raise ValueError(
                        f"The key of the result is not a string or the value is not a RangeModel or NestedRangeModel."
                    )
            return result
        else:
            return None

    def add_attributes(self, **kwargs) -> None:
        """Add attributes to the module.

        Parameters
        ----------
        **kwargs : dict
            The attributes to add to the module.
        """
        for key, value in kwargs.items():
            setattr(self, key, value)

    def execute(self) -> Any:
        if self.local_function_name:
            calculator = get_local_function(
                self,
                self.local_function_name,
                self.pipeline_local_functions_path_first,
                self.debug_mode,
            )
            if calculator is None:
                for port_name, port_data in self.dynamic_ports_out.items():
                    port_data.data = None
                    return None

            result = calculator.run_function()
            if result is not None:
                if not isinstance(result, dict):
                    raise ValueError(
                        f"The result of the function '{self.local_function_name}' (type: {type(result)}) is not a dict."
                    )
                for port_name, port_data in result.items():
                    if port_name in self.dynamic_ports_out.keys():
                        self.dynamic_ports_out[port_name].data = port_data
                    else:
                        raise ValueError(
                            f"The port '{port_name}' from the function result is not found in the current module '{self.mname}'."
                        )
                return result

        for port_name, port_data in self.dynamic_ports_out.items():
            port_data.data = None
        return None


@status_manage
class LoopDataCollector(PipeModule):
    """This widget collects data from input ports across multiple loop iterations
    and outputs the collected data as a list after the loop is completed.

    This is useful in pipeline steps that involve loops, where you want to accumulate
    results from each iteration and then process them together.
    """

    def __init__(
        self,
        mname: str | None = None,
        auto_run: bool = True,
        output_port_type: PortType = PortType.FilesPath,
        output_port_name: str = "OutputCollectedData",
        reset_on_first_run: bool = True,
        max_iterations: int | None = None,
    ) -> None:
        """Initialize LoopDataCollector object.

        Parameters
        ----------
        output_port_type : PortType, default: PortType.FilesPath
            The type of the output port. Can be FilesPath, GeneralArray, etc.

        output_port_name : str, default: "OutputCollectedData"
            The name of the output port.

        reset_on_first_run : bool, default: True
            If True, the collected data will be reset on the first run of each loop cycle.

        max_iterations : int | None, default: None
            Maximum number of iterations to collect. If None, no limit is set.
        """
        super().__init__(mname, auto_run)

        # Setup dynamic input ports to accept any type of data
        self.dynamic_ports_in_type = PortType.General

        # Setup output port with specified type
        self._ports_out = {
            output_port_name: Port(
                output_port_type, pdoc="The collected data from all loop iterations."
            )
        }

        self.output_port_type = output_port_type
        self.output_port_name = output_port_name
        self.reset_on_first_run = reset_on_first_run
        self.max_iterations = max_iterations

        # Internal storage for collected data
        self._collected_data: dict[str, list[Any]] = {}
        self._iteration_count: int = 0
        self._is_first_run: bool = True

    def add_input_port(self, port_name: str, pdoc: str | None = None) -> None:
        """Add a dynamic input port for data collection.

        Parameters
        ----------
        port_name : str
            The name of the input port to add.
        pdoc : str | None, default: None
            The documentation for the port.
        """
        self.add_dynamic_ports_in(port_name, pdoc)
        # Initialize storage for this port
        if port_name not in self._collected_data:
            self._collected_data[port_name] = []

    def reset_collected_data(self) -> None:
        """Reset all collected data."""
        self._collected_data = {}
        self._iteration_count = 0
        # Don't reset _is_first_run flag here since this is called during the first run

        # Reinitialize storage for existing ports
        for port_name in self.dynamic_ports_in.keys():
            self._collected_data[port_name] = []

    def full_reset(self) -> None:
        """Completely reset the collector including the first run flag."""
        self._collected_data = {}
        self._iteration_count = 0
        self._is_first_run = True

        # Reinitialize storage for existing ports
        for port_name in self.dynamic_ports_in.keys():
            self._collected_data[port_name] = []

    def get_collected_data(self) -> dict[str, list[Any]]:
        """Get the currently collected data.

        Returns
        -------
        dict[str, list[Any]]
            Dictionary mapping port names to lists of collected data.
        """
        return self._collected_data.copy()

    @property
    def OutputCollectedData(self) -> Any:
        """Get the collected data from the default output port."""
        return self[self.output_port_name]

    def __getattr__(self, name: str) -> Any:
        """Allow dynamic access to output ports with different names."""
        if name == self.output_port_name:
            return self[self.output_port_name]
        return super().__getattribute__(name)

    def set_cal_params(self, reset: bool = False) -> dict[str, RangeModel] | None:
        """Set calculation parameters. No user parameters needed for this widget."""
        return None

    def execute(self) -> Any:
        """Execute the data collection.

        This method collects data from all input ports and stores them in internal lists.
        The collected data is then output as a list through the output port.
        """
        # Reset data on first run if specified
        if self._is_first_run and self.reset_on_first_run:
            self._collected_data = {}
            self._iteration_count = 0
            # Reinitialize storage for existing ports
            for port_name in self.dynamic_ports_in.keys():
                self._collected_data[port_name] = []
            self._is_first_run = False

        # Collect data from all input ports first
        for port_name, port in self.dynamic_ports_in.items():
            if port.data is not None:
                # Initialize storage if not exists
                if port_name not in self._collected_data:
                    self._collected_data[port_name] = []

                # Add data to collection
                self._collected_data[port_name].append(port.data)

        self._iteration_count += 1

        # Check max iterations limit after processing
        if (
            self.max_iterations is not None
            and self._iteration_count > self.max_iterations
        ):
            # Remove the last added data if we exceeded the limit
            for port_name, port in self.dynamic_ports_in.items():
                if port.data is not None and port_name in self._collected_data:
                    if len(self._collected_data[port_name]) > 0:
                        self._collected_data[port_name].pop()
            self._iteration_count -= 1

        # Prepare output data
        if len(self._collected_data) == 1:
            # If only one input port, output its data list directly
            output_data = list(self._collected_data.values())[0]
        else:
            # If multiple input ports, handle based on output port type
            if self.output_port_type == PortType.FilesPath:
                # For FilesPath, try to flatten all file paths into a single list
                output_data = []
                for port_data_list in self._collected_data.values():
                    for item in port_data_list:
                        if isinstance(item, (list, tuple)):
                            output_data.extend(item)
                        else:
                            output_data.append(item)
            elif self.output_port_type == PortType.GeneralArray:
                # For GeneralArray, convert to list format that's compatible
                # Create a list of dictionaries, one for each iteration
                max_len = (
                    max(len(data_list) for data_list in self._collected_data.values())
                    if self._collected_data
                    else 0
                )
                output_data = []
                for i in range(max_len):
                    iteration_data = {}
                    for port_name, data_list in self._collected_data.items():
                        if i < len(data_list):
                            iteration_data[port_name] = data_list[i]
                    output_data.append(iteration_data)
            else:
                # For other types, keep the dictionary structure
                output_data = dict(self._collected_data)

        # Set output port data
        self._ports_out[self.output_port_name].data = output_data
        return output_data

    def finalize_collection(self) -> Any:
        """Finalize the data collection and return the final result.

        This method can be called explicitly to signal the end of data collection.
        """
        return self.execute()
